"""
PlumeSoft 量化代码库
作者：PlumeSoft
用途：贝叶斯选币炼丹炉
"""

import json
import os
import pickle
import shutil
import time
import pandas as pd
import numpy as np
import joblib
from optuna import Trial
from time import sleep
from datetime import datetime
import program.plumesoft_bys.bys_function as func
from program.plumesoft_bys.bys_back_test import batch_run_stg_back_test
from program.plumesoft_bys.bys_common import (
    add_code_to_stg_file,
    reduce_mem_df,
    dbg,
    log,
)
from program.plumesoft_bys.bys_base import Base_Bayesian
from program import Evaluate as Eva
from program import Config as Cfg
from program import Functions as Fun

import warnings

warnings.filterwarnings("ignore")

################################
# 炼丹炉核心业务函数
################################


class SelectCoin_Bayesian(Base_Bayesian):
    def prepare_global_data(self, params: dict) -> dict:
        """
        准备全局数据，可以在这里进行读取数据、合并数据等数据预处理操作。
        全局数据在主任务中执行一次，被所有后台子任务共享。
        在这个函数中可以使用多任务处理，因为这个函数只在主任务执行。

        Args:
            params (dict): 用来传入和返回公共数据
        Returns:
            dict: 返回修改后的 params 字典
        """

        stime = datetime.now()

        log("炼丹炉开始准备全局数据...")
        dbg(f"基础参数: {params}")

        # 初始化策略和全局参数
        cls, params = func.init_bys_stg_and_params(params)

        # 检查策略是否支持炼丹
        lost_list = [
            name
            for name in [
                "bys_result_user_params",
            ]
            if not hasattr(cls, name)
        ]
        if len(lost_list) > 0:
            log(f"在策略文件中没有找到 {lost_list} 定义，无法支持选币炼丹...")
            exit()

        if params["bys_disable_filter"]:
            log(f"已禁用过滤因子炼丹")
        if params["bys_disable_sort"]:
            log(f"已禁用排序因子炼丹")

        # =====从配置中读取的数据
        # 从Strategy中获取需要回测的factor_list
        factor_list = cls.factor_list
        # 从Strategy中获取需要回测的filter_list
        filter_list = cls.filter_list
        # 从Strategy中获取需要回测持仓周期
        hold_period = cls.hold_period
        # 从Strategy中获取offset
        offset = cls.offset
        # 从Strategy中获取是否使用现货if_use_spot
        if_use_spot = cls.if_use_spot
        # 检查杠杆配置
        Fun.check_leverage(cls, if_use_spot, Cfg.leverage)

        if params["bys_disable_sort"]:
            # ===加载因子信息
            # =根据配置的因子信息加载因子的数据，例如配置了('ILLQ', True, '7', 1)
            # =factor_column_list列表里面会添加  ILLQ_7 然后加载这个字段
            factor_column_list = []
            factor_param_dict = None
            for factor_name, if_reverse, parameter_list, weight in factor_list:
                factor_column_list.append(f"{factor_name}_{str(parameter_list)}")
        else:
            # 加载因子及所有参数的字段
            factor_param_dict, factor_column_list = (
                func.get_factor_params_and_column_list(
                    [factor[0] for factor in factor_list], Cfg.factor_param_list
                )
            )

        if params["bys_disable_filter"]:
            # =根据配置的过滤信息加载过滤的数据，例如配置了('PctChange', 7)
            # =filter_column_list列表里面会添加  PctChange_7 然后加载这个字段
            filter_column_list = []
            filter_param_dict = None
            for factor_name, parameter_list in filter_list:
                filter_column_list.append(f"{factor_name}_{str(parameter_list)}")
        else:
            filter_param_dict, filter_column_list = (
                func.get_factor_params_and_column_list(
                    [factor[0] for factor in filter_list], Cfg.factor_param_list
                )
            )

        # =合并 factor_column_list 和 filter_column_list 的信息（会去重）
        all_factor_list = sorted(list(set(factor_column_list + filter_column_list)))

        # =打印回测信息
        log("炼丹排序列表：", factor_param_dict)
        log("炼丹过滤列表：", filter_param_dict)
        log("持仓周期：", hold_period)
        log("offset：", offset)
        log("是否使用现货：", if_use_spot)

        # ===读取并整理数据
        # =读取之前计算好的pkl数据
        start_time = time.time()
        df = Fun.read_coin(
            Cfg.root_path, hold_period, all_factor_list, if_use_spot, Cfg.n_jobs, offset
        )

        # 持仓时间为1H时，把列表类型的字段转成数值型以节约内存
        if hold_period == "1H":
            for col in ["每小时涨跌幅", "每小时涨跌幅_byclose", "fundingRate"]:
                df[col] = df[col].transform(lambda x: x[0])

        log("读取耗时：", time.time() - start_time)
        if df.empty:
            log("读取df为空，结束当前炼丹")
            exit()

        # ===将数据分为现货数据和合约数据
        # 合约和现货数据分离
        df_spot = df[
            df["symbol_type"] == "spot"
        ]  # 取出现货数据，后续根据配置用来进行过滤和选币
        df_swap = df[
            df["symbol_type"] == "swap"
        ]  # 取出合约数据，后续根据配置用来进行过滤和选币，并替换掉现货中的行情数据

        # 是否使用现货数据进行回测，如果为True，则表明在现货数据中选币；为False，则在合约数据中选币
        if if_use_spot:  # 使用现货数据，则在现货中进行过滤，并选币
            back_test_df = df_spot.copy()
            if back_test_df.empty:  # 如果数据中不存在现货数据就退出
                log("当前为现货模式，但是数据中不存在现货数据，请重新整理数据。")
                exit()
        else:  # 使用现货数据，则在现货中进行过滤，并选币
            back_test_df = df_swap.copy()
        # # 删除选币因子为空的数据
        # back_test_df.dropna(subset=all_factor_list, inplace=True)

        # 基础的字段
        keep_cols = [
            "candle_begin_time",
            "offset",
            "symbol",
            "symbol_type",
            "tag",
            "ret_next",
            "每小时涨跌幅",
            "每小时涨跌幅_byclose",
            "fundingRate",
        ]
        # 策略文件中定义的需要保留的字段
        if hasattr(cls, "user_keep_cols") and cls.user_keep_cols is not None:
            keep_cols.extend(cls.user_keep_cols)

        params["hold_period"] = hold_period
        params["offset"] = offset
        params["factor_list"] = factor_list
        params["factor_param_dict"] = factor_param_dict
        params["filter_list"] = filter_list
        params["filter_param_dict"] = filter_param_dict
        params["keep_cols"] = keep_cols.copy()
        params["all_factor_list"] = all_factor_list.copy()
        params["if_use_spot"] = if_use_spot

        keep_cols.extend(all_factor_list)
        keep_cols = sorted(list(set(keep_cols)))
        log(f"炼丹保留的字段({len(keep_cols)})：{keep_cols}")

        # 只保留一些有用的字段，减小数据大小
        back_test_df = back_test_df[keep_cols]  # 只保留需要的字段

        # 检查过滤因子字段是否存在
        invalid_factors = [
            factor
            for factor in filter_column_list
            if not factor in back_test_df.columns
        ]
        if len(invalid_factors) > 0:
            log(
                f"策略数据集中缺少过滤炼丹配置中定义的以下因子：{invalid_factors}，请检查后再重试，退出程序..."
            )
            exit()

        # 检查排序因子字段是否存在
        invalid_factors = [
            factor
            for factor in factor_column_list
            if not factor in back_test_df.columns
        ]
        if len(invalid_factors) > 0:
            log(
                f"策略数据集中缺少排序炼丹配置中定义的以下因子：{invalid_factors}，请检查后再重试，退出程序..."
            )
            exit()

        # 检查因子字段是否全是空值
        invalid_factors = [
            factor for factor in all_factor_list if back_test_df[factor].isnull().all()
        ]
        if len(invalid_factors) > 0:
            log(
                f"策略数据集中因子 {invalid_factors} 全为空值，可能是数据读取错误，请检查后再重试，退出程序..."
            )
            exit()

        # 检查因子字段空值比例
        cnt = 0
        for factor in all_factor_list:
            na_cnt = back_test_df[factor].isnull().sum()
            if na_cnt / len(back_test_df) > params["factor_null_value_rate_max"]:
                log(
                    f"策略数据集中因子 {factor} 数据的空值占比 {na_cnt / len(back_test_df):.2}，超过允许的最大值 {params['factor_null_value_rate_max']}"
                )
                cnt += 1
        if cnt > 0:
            delay_sec = 30
            log(
                f"警告：策略数据集中存在以上 {cnt} 个因子，数据的空值占比超过允许的最大值 {params['factor_null_value_rate_max']}，这可能导致炼丹效果不稳定。"
            )
            log(
                f"建议优化因子计算或使用 fillna 或者 dropna 对因子数据进行预处理。如果不希望看到该警告，可以在 bys_config.py 中通过 factor_null_value_rate_max 参数来调整该阈值。"
            )
            log(f"延时 {delay_sec} 秒后继续执行...")
            sleep(delay_sec)

        # 内存优化
        if params["reduce_mem_mode"]:
            back_test_df = reduce_mem_df(back_test_df[keep_cols])

        # =构建回测期间，完整的交易时间列表
        # 定义一个开始的基准时间，避免周期转换出现问题
        benchmark = pd.DataFrame(
            pd.date_range(start="2017-01-01", end=params["sample_date_end"], freq="1H")
        )  # 创建2017-01-01至回测结束时间的1H列表
        benchmark.rename(columns={0: "candle_begin_time"}, inplace=True)
        trading_time_list = Fun.create_trading_time(
            benchmark,
            hold_period,
            offset,
            params["sample_date_start"],
            (
                params["sample_date_end"]
                if params["sample_date_end"] is not None
                else benchmark["candle_begin_time"].max()
            ),
        )

        # 创造空的时间周期表，用于填充不选币的周期
        empty_df = Fun.create_empty_data(benchmark, hold_period, offset=offset)

        # 参数很大时使用交换文件提升性能
        if params["study_params_swap_pkl"] is not None:
            if params["use_joblib_sharemem"]:
                joblib.dump(
                    [back_test_df, df_swap, benchmark, trading_time_list, empty_df],
                    params["study_params_swap_pkl"],
                )
            else:
                pickle.dump(
                    [back_test_df, df_swap, benchmark, trading_time_list, empty_df],
                    open(params["study_params_swap_pkl"], "wb"),
                )
        else:
            params["back_test_df"] = back_test_df
            params["df_swap"] = df_swap
            params["benchmark"] = benchmark
            params["trading_time_list"] = trading_time_list
            params["empty_df"] = empty_df

        log(f"准备全局数据完成，总耗时：{datetime.now() - stime}")

        return params

    def prepare_task_data(self, params: dict) -> dict:
        """
        准备任务用到的数据，可以在这里执行数据读取、初始化等操作。
        这个函数只在每个后台任务启动时执行一次。
        注意：尽量不要在这个函数中使用多任务处理，因为框架本身已经是多任务的了。

        Args:
            params (dict): 用来传入和返回公共数据
        Returns:
            dict: 返回修改后的 params 字典
        """

        log("炼丹炉开始准备任务数据...")

        # =====动态导入选币策略
        cls = __import__(
            f'program.{params["base_stg_folder"]}.{params["base_stg_file"]}',
            fromlist=("",),
        )

        # 从交换文件中读取大数据
        if params["study_params_swap_pkl"] is not None:
            if params["use_joblib_sharemem"]:
                # 如果打开了 joblib 的 sharemem 模式，使用 joblib 的 load 函数使用共享内存节约内存占用
                (
                    params["back_test_df"],
                    params["df_swap"],
                    params["benchmark"],
                    params["trading_time_list"],
                    params["empty_df"],
                ) = joblib.load(params["study_params_swap_pkl"], mmap_mode="r")
            else:
                (
                    params["back_test_df"],
                    params["df_swap"],
                    params["benchmark"],
                    params["trading_time_list"],
                    params["empty_df"],
                ) = pickle.load(open(params["study_params_swap_pkl"], "rb"))

        # 加载任务公共参数
        params["cls"] = cls

        return params

    def execute_trial(
        self, trial: Trial, params: dict, result_queue, score_queue
    ) -> float:
        """
        炼丹迭代任务回测主函数。
        这个函数会执行多次，每次生成新的参数，进行回测，返回评分值。
        注意：尽量不要在这个函数中使用多任务处理，因为框架本身已经是多任务的了。

        Args:
            trial (Trial): 测试对象，用于生成炼丹参数
            params (dict): 用于炼丹的配置和公共参数，包含由 prepare_data 提供的数据
            result_queue (queue): 用于向主任务返回单次测试的结果，DataFrame 格式，其中需要包含 score 评分字段
            score_queue (queue): 用于向主任务返回高分数据和参数，包含4个变量

        Returns:
            float: Optuna 要求的评分值，越高表示测试结果越好
        """
        # ret_dt 是字典类型，用来保存测试结果，里面可以包含任何需要保存到测试结果 csv 文件中的字段，必须有一个是 score 得分字段
        ret_dt = {"迭代": trial.number}
        id_str = f"迭代 {trial.number}"

        start_time = time.time()

        # 取得测试使用的数据
        cls = params["cls"]

        back_test_df = params["back_test_df"]
        df_swap = params["df_swap"]
        benchmark = params["benchmark"]
        trading_time_list = params["trading_time_list"]
        _empty_df = params["empty_df"]

        hold_period = params["hold_period"]
        offset = params["offset"]
        factor_list = params["factor_list"]
        factor_param_dict = params["factor_param_dict"]
        filter_list = params["filter_list"]
        filter_param_dict = params["filter_param_dict"]
        bys_eval_score_dict = params["bys_eval_score_dict"]
        keep_cols = params["keep_cols"]
        if_use_spot = params["if_use_spot"]

        if params["bys_disable_filter"]:
            # 禁用过滤炼丹
            filter_factor_list = filter_list
        else:
            # 生成过滤炼丹参数
            filter_factor_list = func.suggest_filter_list(
                trial, params, "filter", filter_param_dict, "过滤"
            )
        ret_dt["过滤参数"] = filter_factor_list
        dbg(f"{id_str} 过滤参数: ", filter_factor_list)
        if len(factor_list) > 0 and len(filter_factor_list) == 0:
            log(f"{id_str} 过滤因子列表为空，放弃回测。")
            return -2000

        # 配置了因子互斥模式时，需要排除过滤因子
        factor_exclude_list = filter_factor_list if params["factor_mutex_mode"] else []

        if params["bys_disable_sort"]:
            # 禁用排序炼丹
            sort_factor_list = factor_list
        else:
            # 生成排序因子
            sort_factor_list = func.suggest_sort_factor_list(
                trial, params, factor_param_dict, factor_exclude_list
            )
        ret_dt["排序参数"] = sort_factor_list
        dbg(f"{id_str} 排序参数: ", sort_factor_list)
        if len(sort_factor_list) == 0:
            log(f"{id_str} 排序因子列表为空，放弃回测。")
            return -2000

        sort_column_list = [
            factor + "_" + str(param) for factor, _, param, _ in sort_factor_list
        ]
        filter_column_list = [
            factor + "_" + str(param) for factor, param in filter_factor_list
        ]

        # 定义过滤和排序因子的宏对应的因子列表
        macro_dict = {"$过滤因子$": filter_column_list, "$排序因子$": sort_column_list}
        # 生成用户参数
        user_params = func.suggest_user_params(
            cls,
            trial,
            params["bys_user_param_dict"],
            filter_column_list,
            sort_column_list,
            macro_dict=macro_dict,
        )
        ret_dt["用户参数"] = user_params

        # ===计算因子
        # 调用Strategy.calc_factor计算最终用来选币的因子，具体的因子计算方法可以在Strategy.calc_factor中自定义
        setattr(cls, "factor_list", sort_factor_list)
        setattr(cls, "filter_list", filter_factor_list)
        setattr(cls, "bys_result_user_params", user_params)
        back_test_df = cls.calc_factor(
            back_test_df[
                list(set(keep_cols + sort_column_list + filter_column_list))
            ].copy()
        )

        # =计算好因子后的一些数据整理工作
        # 获取做多和做空因子的字段名，本框架支持多头和空头是不同的因子
        long_short_columns = list({cls.long_factor, cls.short_factor})
        # 后期需要保存的列 = 多空因子字段 + 过滤因子字段
        save_columns = list(set(long_short_columns + filter_column_list))
        # 只保留一些有用的字段，减小数据大小
        back_test_df = back_test_df[keep_cols + save_columns]  # 只保留需要的字段

        # ===对回测数据进行过滤
        # 过滤的具体方法可以在Strategy.before_filter中自定义，故多空可能是分开过滤的，所以需要多头给一个df，空头也给一个df
        long_df, short_df = cls.before_filter(back_test_df)
        # 如果使用现货数据，多头选币不受影响，但是空头选币只能在有合约的现货中进行选币
        if if_use_spot:
            short_df = short_df[short_df["tag"] == "HasSwap"]  # 保留有合约的现货

        # ===选币操作
        # 根据选币数量进行多空选币操作
        select_coin = Fun.select_long_and_short_coin(
            long_df,
            short_df,
            cls.long_select_coin_num,
            cls.short_select_coin_num,
            long_factor=cls.long_factor,
            short_factor=cls.short_factor,
        )
        # 如果是现货模式，将现货涨跌幅数据换成合约涨跌幅数据
        if if_use_spot:
            select_coin = Fun.transfer_swap(
                select_coin, df_swap.copy(), Cfg.special_symbol_dict
            )

        # ===计算资金曲线前的一些操作
        # 针对合约和现货给予不同的手续费率
        select_coin.loc[select_coin["symbol_type"] == "spot", "rate"] = Cfg.spot_c_rate
        select_coin.loc[select_coin["symbol_type"] == "swap", "rate"] = Cfg.swap_c_rate
        select_coin["lvg"] = Cfg.leverage

        # 保留指定字段
        select_coin = select_coin[
            [
                "candle_begin_time",
                "offset",
                "symbol",
                "symbol_type",
                "方向",
                "每小时涨跌幅",
                "每小时涨跌幅_byclose",
                "fundingRate",
                "rate",
                "lvg",
                "ret_next",
            ]
        ]

        # 在选币信息上带上symbol和方向的信息，例如 做空BTC-USDT 会变成 BTC-USDT(-1)
        select_coin["选币"] = (
            select_coin["symbol"]
            + "("
            + select_coin["symbol_type"]
            + ","
            + select_coin["方向"].astype(str)
            + ")"
            + " "
        )

        # 把选币结果重新排序以及重置索引
        select_coin.sort_values(by="candle_begin_time", inplace=True)
        select_coin.reset_index(drop=True, inplace=True)

        # 持仓时间为1H时，计算回测前把数值型字段再转回列表
        if hold_period == "1H":
            for col in ["每小时涨跌幅", "每小时涨跌幅_byclose", "fundingRate"]:
                select_coin[col] = select_coin[col].transform(lambda x: [x])

        dbg(f"{id_str} 选币完成：", time.time() - start_time)

        # ===计算资金曲线
        # =将多头和空头拆开计算
        long_df = select_coin[select_coin["方向"] == 1]
        short_df = select_coin[select_coin["方向"] == -1]

        # 调用函数计算扣除资金费 & 手续费后的每日净值
        long_df = Fun.cal_net_value(long_df, hold_period, trading_time_list)
        short_df = Fun.cal_net_value(short_df, hold_period, trading_time_list)

        # =计算合约仓位信息
        # PS: 放在这里计算，是因为保留最原始的涨跌幅数据，我们需要单独计算合约仓位的净值
        swap_net_value = Fun.calc_swap_pos_net_value(
            select_coin, Cfg.leverage, hold_period, Cfg.margin_rate, trading_time_list
        )
        dbg(f"{id_str} 计算多空资金曲线 + 合约仓位资金曲线：", time.time() - start_time)

        # 填充缺失数据
        empty_df = _empty_df.copy()
        empty_df.update(swap_net_value)
        swap_net_value = empty_df.copy()
        swap_net_value["选币"].fillna(method="ffill", inplace=True)
        swap_net_value.dropna(subset=["选币"], inplace=True)
        swap_net_value["每小时资金曲线"] = swap_net_value["每小时资金曲线"].apply(list)

        # 将多空方向处理后的数据合并起来
        select_coin = pd.concat([long_df, short_df], ignore_index=True)
        dbg(f"{id_str} 多空资金曲线：", time.time() - start_time)

        # ===计算当周期资金曲线
        # 将多空选币的资金曲线合并成当周期内的资金曲线
        group = select_coin.groupby("candle_begin_time")
        # 存储周期内的资金曲线
        merge_df = pd.DataFrame()
        merge_df["选币"] = group["选币"].sum()
        merge_df["offset"] = group["offset"].last()
        # 将多空方向的资金曲线合并起来
        merge_df["每小时资金曲线"] = group["每小时资金曲线"].apply(
            lambda x: np.array(x).mean(axis=0)
        )
        # 计算周期的涨跌幅
        merge_df["周期涨跌幅"] = merge_df["每小时资金曲线"].apply(lambda x: x[-1] - 1)
        # 通过每小时资金曲线，计算每小时账户资金的涨跌幅
        merge_df["每小时涨跌幅"] = merge_df["每小时资金曲线"].apply(
            lambda x: list(np.array(x) / np.array([1.0] + list(x[:-1])) - 1)
        )
        # 计算整个周期多空的平均调仓
        merge_df["多空调仓比例"] = group["调仓比例"].mean()

        # 将以close计算的每小时涨跌幅(用来轮动)也做相同的操作
        merge_df["每小时资金曲线_byclose"] = group["每小时资金曲线_byclose"].apply(
            lambda x: np.array(x).mean(axis=0)
        )
        merge_df["周期涨跌幅_byclose"] = merge_df["每小时资金曲线_byclose"].apply(
            lambda x: x[-1] - 1
        )
        merge_df["每小时涨跌幅_byclose"] = merge_df["每小时资金曲线_byclose"].apply(
            lambda x: list(np.array(x) / np.array([1.0] + list(x[:-1])) - 1)
        )

        # 判断是否生成多头、空头的资金曲线。
        if Cfg.if_all_curves:
            # 计算多空方向的资金曲线
            if long_df.empty:
                merge_df["多头每小时资金曲线"] = None
                merge_df["多头周期涨跌幅"] = None
                merge_df["多头每小时涨跌幅"] = None
                merge_df["多头调仓比例"] = None

                # 以close计算的涨跌幅也做相同的操作
                merge_df["多头每小时资金曲线_byclose"] = None
                merge_df["多头周期涨跌幅_byclose"] = None
                merge_df["多头每小时涨跌幅_byclose"] = None
            else:
                long_df = long_df.set_index("candle_begin_time")
                merge_df["多头每小时资金曲线"] = long_df.loc[:, "每小时资金曲线"]
                merge_df["多头周期涨跌幅"] = merge_df["多头每小时资金曲线"].apply(
                    lambda x: x[-1] - 1
                )
                merge_df["多头每小时涨跌幅"] = merge_df["多头每小时资金曲线"].apply(
                    lambda x: list(np.array(x) / np.array([1.0] + list(x[:-1])) - 1)
                )
                merge_df["多头调仓比例"] = long_df.loc[:, "调仓比例"]

                # 同上
                merge_df["多头每小时资金曲线_byclose"] = long_df.loc[
                    :, "每小时资金曲线_byclose"
                ]
                merge_df["多头周期涨跌幅_byclose"] = merge_df[
                    "多头每小时资金曲线_byclose"
                ].apply(lambda x: x[-1] - 1)
                merge_df["多头每小时涨跌幅_byclose"] = merge_df[
                    "多头每小时资金曲线_byclose"
                ].apply(
                    lambda x: list(np.array(x) / np.array([1.0] + list(x[:-1])) - 1)
                )

            if short_df.empty:
                merge_df["空头每小时资金曲线"] = None
                merge_df["空头周期涨跌幅"] = None
                merge_df["空头每小时涨跌幅"] = None
                merge_df["空头调仓比例"] = None

                # 同上
                merge_df["空头每小时资金曲线_byclose"] = None
                merge_df["空头周期涨跌幅_byclose"] = None
                merge_df["空头每小时涨跌幅_byclose"] = None
            else:
                short_df = short_df.set_index("candle_begin_time")
                merge_df["空头每小时资金曲线"] = short_df.loc[:, "每小时资金曲线"]
                merge_df["空头周期涨跌幅"] = merge_df["空头每小时资金曲线"].apply(
                    lambda x: x[-1] - 1
                )  # 计算周期的涨跌幅
                merge_df["空头每小时涨跌幅"] = merge_df["空头每小时资金曲线"].apply(
                    lambda x: list(np.array(x) / np.array([1.0] + list(x[:-1])) - 1)
                )  # 通过每小时资金曲线，计算每小时账户资金的涨跌幅
                merge_df["空头调仓比例"] = short_df.loc[:, "调仓比例"]

                # 同上
                merge_df["空头每小时资金曲线_byclose"] = short_df.loc[
                    :, "每小时资金曲线_byclose"
                ]
                merge_df["空头周期涨跌幅_byclose"] = merge_df[
                    "空头每小时资金曲线_byclose"
                ].apply(lambda x: x[-1] - 1)
                merge_df["空头每小时涨跌幅_byclose"] = merge_df[
                    "空头每小时资金曲线_byclose"
                ].apply(
                    lambda x: list(np.array(x) / np.array([1.0] + list(x[:-1])) - 1)
                )

        # =====计算小时级别的资金曲线
        # 创造空的时间周期表，用于填充不选币的周期
        empty_df = _empty_df.copy()
        # 填充缺失数据
        empty_df.update(merge_df)
        merge_df = empty_df.copy()
        merge_df["选币"].fillna(method="ffill", inplace=True)
        merge_df.dropna(subset=["选币"], inplace=True)
        dbg(f"{id_str} 填充当周期资金曲线：", time.time() - start_time)

        # 将合并后的资金曲线，与benchmark合并
        if Cfg.if_all_curves:  # 如果生成多头、空头的资金曲线，在merge时多加两列
            equity = pd.merge(
                left=benchmark,
                right=merge_df[
                    ["选币", "多空调仓比例", "多头调仓比例", "空头调仓比例"]
                ],
                on=["candle_begin_time"],
                how="left",
                sort=True,
            )
        else:
            equity = pd.merge(
                left=benchmark,
                right=merge_df[["选币", "多空调仓比例"]],
                on=["candle_begin_time"],
                how="left",
                sort=True,
            )
        # 填充选币数据
        equity["选币"].fillna(method="ffill", inplace=True)
        equity.dropna(subset=["选币"], inplace=True)
        # 将每小时涨跌幅数据，填充到index中
        equity["涨跌幅"] = Fun.revise_data_length(
            merge_df["每小时涨跌幅"].sum(), len(equity)
        )
        equity["涨跌幅"].fillna(value=0, inplace=True)
        # 计算最终每小时净值变化
        equity["多空资金曲线"] = (equity["涨跌幅"] + 1).cumprod()
        equity["多空资金曲线"].fillna(value=1, inplace=True)
        # 计算爆仓
        if not swap_net_value.empty:
            equity["合约资金曲线"] = Fun.revise_data_length(
                swap_net_value["每小时资金曲线"].sum(), len(equity), value=1
            )
            equity.loc[equity["合约资金曲线"] == 0, "是否爆仓"] = 1
            equity.loc[equity["多空资金曲线"] < Cfg.margin_rate, "是否爆仓"] = 1
        else:
            equity["是否爆仓"] = 0
        # 这里计算合约仓位爆仓之后，直接将现货仓位也同步归零，实际实盘的时候的会留有部分现货仓位的底仓
        equity["是否爆仓"].fillna(method="ffill", inplace=True)
        equity.loc[equity["是否爆仓"] == 1, "多空资金曲线"] = 0
        equity.loc[equity["是否爆仓"] == 1, "涨跌幅"] = 0
        equity.loc[equity["是否爆仓"] == 1, "多空调仓比例"] = 0
        # 保存资金曲线文件
        equity["本周期多空涨跌幅"] = equity["涨跌幅"]

        # 同上
        equity["涨跌幅_byclose"] = Fun.revise_data_length(
            merge_df["每小时涨跌幅_byclose"].sum(), len(equity)
        )
        equity["涨跌幅_byclose"].fillna(value=0, inplace=True)
        equity["多空资金曲线_byclose"] = (equity["涨跌幅_byclose"] + 1).cumprod()
        equity["多空资金曲线_byclose"].fillna(value=1, inplace=True)
        equity.loc[equity["是否爆仓"] == 1, "多空资金曲线_byclose"] = 0
        equity.loc[equity["是否爆仓"] == 1, "涨跌幅_byclose"] = 0

        # 计算多空策略的评价指标
        rtn, year_return, month_return = Eva.strategy_evaluate(
            equity,
            net_col="多空资金曲线",
            pct_col="本周期多空涨跌幅",
            turnover_col="多空调仓比例",
        )

        # 如果生乘多头、空头的资金曲线
        if Cfg.if_all_curves:
            # 合并多头、空头每小时涨跌幅 并 计算资金曲线
            equity["多头涨跌幅"] = Fun.revise_data_length(
                merge_df["多头每小时涨跌幅"].sum(), len(equity)
            )
            equity["多头涨跌幅"].fillna(value=0, inplace=True)
            equity["空头涨跌幅"] = Fun.revise_data_length(
                merge_df["空头每小时涨跌幅"].sum(), len(equity)
            )
            equity["空头涨跌幅"].fillna(value=0, inplace=True)
            equity["多头资金曲线"] = (equity["多头涨跌幅"] + 1).cumprod()
            equity["多头资金曲线"].fillna(value=1, inplace=True)
            equity["空头资金曲线"] = (equity["空头涨跌幅"] + 1).cumprod()
            equity["空头资金曲线"].fillna(value=1, inplace=True)

            # 同上
            equity["多头涨跌幅_byclose"] = Fun.revise_data_length(
                merge_df["多头每小时涨跌幅_byclose"].sum(), len(equity)
            )
            equity["多头涨跌幅_byclose"].fillna(value=0, inplace=True)
            equity["多头资金曲线_byclose"] = (
                equity["多头涨跌幅_byclose"] + 1
            ).cumprod()
            equity["多头资金曲线_byclose"].fillna(value=1, inplace=True)
            equity["空头涨跌幅_byclose"] = Fun.revise_data_length(
                merge_df["空头每小时涨跌幅_byclose"].sum(), len(equity)
            )
            equity["空头涨跌幅_byclose"].fillna(value=0, inplace=True)
            equity["空头资金曲线_byclose"] = (
                equity["空头涨跌幅_byclose"] + 1
            ).cumprod()
            equity["空头资金曲线_byclose"].fillna(value=1, inplace=True)

            # =====策略评价
            # 计算纯多头
            long_rtn, long_year_return, long_month_return = Eva.strategy_evaluate(
                equity,
                net_col="多头资金曲线",
                pct_col="多头涨跌幅",
                turnover_col="多头调仓比例",
            )
            # 计算纯空头
            short_rtn, short_year_return, short_month_return = Eva.strategy_evaluate(
                equity,
                net_col="空头资金曲线",
                pct_col="空头涨跌幅",
                turnover_col="空头调仓比例",
            )

            # 多空策略评价合并
            rtn_all = pd.concat([rtn.T, long_rtn.T, short_rtn.T], ignore_index=True)
            rtn_all = rtn_all.T.rename(columns={0: "中性", 1: "多头", 2: "空头"})
            dbg(f"\n\n{id_str} ", rtn_all)
            # 多空分年收益率
            year_return_all = pd.concat(
                [year_return.T, long_year_return.T, short_year_return.T],
                ignore_index=True,
            )
            year_return_all = year_return_all.T.rename(
                columns={0: "中性", 1: "多头", 2: "空头"}
            )
            dbg(f"\n\n{id_str} 分年收益率：\n", year_return_all)
            # 分月双分组
            dbg(
                f"\n\n{id_str} 多头分月总收益率：\n",
                long_month_return,
                f"\n\n{id_str} 空头分月总收益率：\n",
                short_month_return,
                f"\n\n{id_str} 中性分月总收益率：\n",
                month_return,
            )
        else:
            dbg(f"{id_str} 策略评价：", time.time() - start_time)
            dbg(
                rtn,
                f"\n\n{id_str} 分年收益率：\n",
                year_return,
                f"\n\n{id_str} 分月收益率：\n",
                month_return,
            )

        # 计算综合评分
        score_sum = 0
        equity_zx = equity[["candle_begin_time", "多空资金曲线"]].rename(columns={"多空资金曲线": "净值"})
        equity_long = equity[["candle_begin_time", "多头资金曲线"]].rename(columns={"多头资金曲线": "净值"})
        equity_short = equity[["candle_begin_time", "空头资金曲线"]].rename(columns={"空头资金曲线": "净值"})
        core_dt = {}
        data_dt = {
            "中性": [equity_zx, rtn.T, year_return, month_return],
            "多头": [equity_long, long_rtn.T, long_year_return, long_month_return],
            "空头": [equity_short, short_rtn.T, short_year_return, short_month_return],
        }
        if set(list(data_dt.keys())) & set(list(bys_eval_score_dict.keys())):
            # 评价表中有对应类型的评分
            eval_dt = bys_eval_score_dict
        else:
            # 兼容早期版本没有区分不同类型的评分
            eval_dt = {"中性": bys_eval_score_dict}
        for tp, data_list in data_dt.items():
            if tp in eval_dt.keys():
                # 计算评分
                score_select, r_dt, c_dt = func.cal_rtn_score(
                    params,
                    cls,
                    tp,
                    eval_dt[tp],
                    data_list[0],
                    data_list[1],
                    data_list[2],
                    data_list[3],
                    is_sample=True,
                )
                score_sum += score_select
                core_dt[f"{tp}_回测结果"] = r_dt
                core_dt[f"{tp}_评分明细"] = c_dt
                for key, val in r_dt.items():
                    ret_dt[f"{key}_{tp}"] = val
                for key, val in c_dt.items():
                    ret_dt[f"{key}_{tp}"] = val

        # ret_dt 中需要包含 score 字段
        ret_dt["回测数据"] = str(core_dt)
        ret_dt["score"] = score_sum

        # 设置参数到trial中
        for key in ["过滤参数", "排序参数", "用户参数", "回测数据"]:
            if ret_dt.get(key, None):
                trial.set_user_attr(key, str(ret_dt[key]))

        # 将 ret_dt 转换为 DataFrame 格式，通过 result_queue.put() 将测试结果保存到主任务中
        result_queue.put(pd.DataFrame([ret_dt]))

        # 保存高分记录
        score_queue.put([trial.number, score_sum, str(ret_dt), str(core_dt)])

        # 返回评分值，用于训练任务评估训练结果和方向，分数越大越好
        return score_sum

    def output_result_stg_files(
        self,
        params,
        result_df,
        top_num,
        back_test_num,
        suffix="炼丹排名",
        open_pic=True,
    ):
        log(f"正在输出炼丹得分前 {top_num} 的策略文件...")
        result_df = (
            result_df.sort_values("score", ascending=False)
            .drop_duplicates(subset=["回测数据"], keep="first", ignore_index=True)
            .head(top_num)
        )
        stg_file_list = []
        for n, rec in enumerate(result_df.to_records()):
            stg_file = f'{params["base_stg_file"]}${suffix}${n+1:03}'
            stg_file_list.append(stg_file)
            src_fn = os.path.join(
                Cfg.root_path,
                "program",
                params["base_stg_folder"],
                params["base_stg_file"] + ".py",
            )
            dst_fn = os.path.join(
                Cfg.root_path, "program", params["output_stg_folder"], f"{stg_file}.py"
            )
            os.makedirs(os.path.dirname(dst_fn), exist_ok=True)
            code = f"\n# 炼丹时间: {datetime.now()}\n"
            code += f'# 基础策略: {params["base_stg_file"]}\n'
            round_id = rec["轮次"] if "轮次" in result_df.columns else params.get("轮次", None)
            if round_id is not None:
                code += f'# 批量轮次: {round_id}\n'
            code += f'# 迭代次数: {rec["迭代"]}\n'
            code += f'# 炼丹得分: {rec["score"]} 排名: {n+1}\n'

            for col in ["回测数据", "样本外回测数据"]:
                if col in result_df.columns:
                    str_res = func.dict_to_json_str(rec[col])
                    str_res = str_res.replace("\n", "\n# ")
                    code += f"# {col}: {str_res}\n\n"
            if not params["bys_disable_sort"]:
                str_sort = func.list_to_code(rec["排序参数"])
                code += f"factor_list = {str_sort}\n\n"
            if not params["bys_disable_filter"]:
                str_filter = func.list_to_code(rec["过滤参数"])
                code += f"filter_list = {str_filter}\n\n"
            str_user = func.dict_to_json_str(rec["用户参数"])
            code += f"bys_result_user_params = {str_user}\n\n"
            add_code_to_stg_file(
                src_fn,
                dst_fn,
                "# @$开始定义选币炼丹结果数据$@#",
                "# @$结束定义选币炼丹结果数据$@#",
                code,
                append_mode=False,
            )
        log(
            f'已生成 {len(stg_file_list)} 个策略文件，保存在 {params["output_stg_folder"]} 文件夹中。'
        )
        if back_test_num > 0:
            log(
                f"开始使用框架 Config.py 中配置的参数执行策略回测，回测区间：{Cfg.start_date} - {Cfg.end_date}"
            )
            batch_run_stg_back_test(
                params["output_stg_folder"],
                stg_file_list[0:back_test_num],
                Cfg.start_date,
                Cfg.end_date,
                if_plot=True,
                n_jobs=params["stg_back_test_n_jobs"],
                pic_dir=params["pic_path"],
                open_pic=open_pic,
            )
        return stg_file_list

    def study_finished(
        self, params: dict, result_df: pd.DataFrame, study_df: pd.DataFrame
    ):
        log(
            f"学习完成，总共运行 {len(study_df)} 次迭代，其中 {len(result_df)} 组有效测试，占比 {(len(result_df)/len(study_df)*100):.1f}%"
        )

        stg_file_list = []
        if not "disable_output_stg" in params.keys():
            stg_file_list = self.output_result_stg_files(
                params,
                result_df,
                params["output_stg_top_num"],
                params["output_stg_back_test_num"],
            )

        # 移动工作目录下的文件到完成目录
        func.move_files_to_finished_path(params, stg_file_list)
